#!/usr/bin/env python
# -*- coding: UTF-8 -*-
'''
@Project ：V2 
@File    ：gas_pipeline_schedule1.py
@IDE     ：PyCharm 
@Author  ：郭星
@Date    ：2025/9/23 14:33
再加一条，如果这瓶气有多个组分是要一个一个充的不能直接一下子充好，无论是分批还是单独充都是要一个组分一个组分充气的
'''
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Gas Mixing -> QC Scheduling (component-wise MIX; batched; 1-pass or serial QC)
-------------------------------------------------------------------------------
- MIX: same recipe (components+concentrations) bottles can be batched,
       BUT charging is done component-by-component in a specified order.
       For a batch with recipe {c1,c2,...}, we schedule sub-batches:
         Charge(c1) -> Charge(c2) -> ...  (precedence within the batch)
       Parallel mixers; sequence-dependent setup by last_component.
- QC : each QC device measures exactly ONE component.
       Per bottle:
         * 1-pass: occupy all needed devices simultaneously (same [st,en]),
                   duration = max(single-device times); NO mid-detach.
         * serial: test devices one-by-one (sum).
- Gantt: MIX step bars labeled with component; QC 1-pass bars same color+label.

Author: you
"""

import random
import numpy as np
from dataclasses import dataclass
from typing import List, Dict, Tuple, Set, Optional

try:
    import matplotlib.pyplot as plt
    HAS_MPL = True
except Exception:
    HAS_MPL = False


# ============================ CONFIG ============================ #

CONFIG = {
    # ---- scale ----
    "NUM_BOTTLES": 20,
    "NUM_MIX_MACHINES": 2,
    "NUM_QC_DEVICES": 6,
    "MAX_COMPONENTS_PER_BOTTLE": 3,

    # ---- MIX (charging) ----
    "MIX_BATCH_CAP": 300,                 # max bottles per batch (same recipe key)
    "MIX_RATE_PER_H": [180.0, 150.0],     # bottles/hour per mixer
    "MIX_SETUP_SAME": 0.05,               # if last_component == current_component
    "MIX_SETUP_DIFF": 0.35,               # if last_component != current_component
    # 每个组分一瓶的充气时间（小时）
    "COMPONENT_FILL_TIME": {
        "O2": 0.015, "CO2": 0.020, "CH4": 0.022, "CO": 0.018, "H2": 0.012, "N2": 0.010, "C2H6": 0.025
    },
    # 充气顺序优先级（左靠前）。不在此表的按字母序最后排。
    "CHARGE_PRIORITY": ["CO2", "O2", "CH4", "H2", "CO", "N2", "C2H6"],

    # ---- QC (testing) ----
    "COMPONENT_TYPES": ["O2", "CO2", "CH4", "CO", "H2", "N2", "C2H6"],  # available types (cycled over devices)
    "COMPONENT_TEST_TIME": {  # per bottle single-device time (hours)
        "O2": 0.18, "CO2": 0.22, "CH4": 0.25, "CO": 0.20, "H2": 0.15, "N2": 0.12, "C2H6": 0.28
    },
    "QC_SETUP_TIME": 0.12,     # different signature (component) switch

    # ---- policy: decide 1-pass vs serial for each bottle ----
    #   "always_1pass", "always_serial", "by_ratio", "by_compare"
    "QC_MODE_POLICY": "by_compare",
    "ONEPASS_RATIO": 0.6,      # if max(single) <= ratio*sum(single) -> 1-pass
    "ONEPASS_RANDOM_RATIO": 0.7,

    # ---- plot ----
    "LABEL_WRAP_CHARS": 18,

    "RANDOM_SEED": 7
}


# ============================ Data Models ============================ #

@dataclass(frozen=True)
class Recipe:
    components: Tuple[Tuple[str, float], ...]  # (component, concentration%)

    def key(self) -> Tuple[Tuple[str,float], ...]:
        # 完全相同的 组分+浓度 视为同配方，可合批
        return tuple(sorted(self.components, key=lambda x: x[0]))


@dataclass
class Bottle:
    id: int
    recipe: Recipe


@dataclass
class MixTask:
    start: float
    end: float
    machine: int
    label: str   # e.g., "Batch(S1,S3)-CO2"


@dataclass
class QCTask:
    start: float
    end: float
    device: int
    label: str
    group_key: Optional[Tuple[int,float,float]] = None  # for 1-pass group coloring


# ============================ Random Instance ============================ #

class RandomData:
    def __init__(self, cfg):
        self.cfg = cfg
        random.seed(cfg["RANDOM_SEED"])
        np.random.seed(cfg["RANDOM_SEED"])

    def gen(self):
        N = self.cfg["NUM_BOTTLES"]
        D = self.cfg["NUM_QC_DEVICES"]
        comps_all = self.cfg["COMPONENT_TYPES"]
        Kmax = self.cfg["MAX_COMPONENTS_PER_BOTTLE"]

        # map device -> single component type (cycled)
        dev_to_comp: List[str] = [comps_all[i % len(comps_all)] for i in range(D)]

        # bottles: choose 1..Kmax components that are testable by some device
        uniq_types = list(dict.fromkeys(dev_to_comp))  # keep order
        bottles: List[Bottle] = []
        for i in range(N):
            k = np.random.randint(1, Kmax+1)
            comps = random.sample(uniq_types, k)
            # concentrations (demo only)
            comps_t = tuple((c, round(np.random.uniform(0.1,5.0),3)) for c in comps)
            bottles.append(Bottle(i, Recipe(comps_t)))
        return bottles, dev_to_comp


# ============================ Scheduler ============================ #

class Scheduler:
    def __init__(self, cfg, bottles: List[Bottle], dev_to_comp: List[str]):
        self.cfg = cfg
        self.bottles = bottles
        self.N = len(bottles)
        self.dev_to_comp = dev_to_comp
        self.D_mix = cfg["NUM_MIX_MACHINES"]
        self.D_qc = cfg["NUM_QC_DEVICES"]

        rates = cfg["MIX_RATE_PER_H"]
        if len(rates) < self.D_mix:
            mul = (self.D_mix + len(rates) - 1) // len(rates)
            rates = (rates * mul)[:self.D_mix]
        self.mix_rates = rates

    # ---------- helper: component order for charging ----------
    def charging_sequence(self, comps: List[str]) -> List[str]:
        pri = {c: idx for idx, c in enumerate(self.cfg["CHARGE_PRIORITY"])}
        def key(c):
            return (pri.get(c, 10_000), c)  # others after, alphabetical
        return sorted(comps, key=key)

    # ---------- MIX: component-by-component, with batching ----------
    def schedule_mixing(self) -> Tuple[Dict[int,float], List[MixTask]]:
        """
        Build batches by exact recipe key; for each batch,
        execute sub-batches Charge(component) in the configured order.
        Each sub-batch time = (#bottles_in_batch) * FILL_TIME[component] at chosen mixer rate
                              (here we use time directly per bottle, independent of mixer rate);
        Setup depends on last_component on that mixer.
        """
        cap = self.cfg["MIX_BATCH_CAP"]
        setup_same = self.cfg["MIX_SETUP_SAME"]
        setup_diff = self.cfg["MIX_SETUP_DIFF"]
        fill_t = self.cfg["COMPONENT_FILL_TIME"]

        # group bottles by recipe key
        key2ids: Dict[Tuple, List[int]] = {}
        for b in self.bottles:
            key2ids.setdefault(b.recipe.key(), []).append(b.id)

        # split each group into batches (<=cap)
        batches: List[List[int]] = []
        for _, ids in key2ids.items():
            cur=[]
            for bid in ids:
                cur.append(bid)
                if len(cur) >= cap:
                    batches.append(cur); cur=[]
            if cur: batches.append(cur)

        # a stable priority: larger batches first
        batches.sort(key=lambda bl: -len(bl))

        # per-mixer ready time and last charged component
        ready = [0.0]*self.D_mix
        last_comp: List[Optional[str]] = [None]*self.D_mix

        # for each batch we must do component sequence one by one
        mix_tasks: List[MixTask] = []
        batch_last_end: Dict[int,float] = {}  # bottle -> end time of final component

        for bl in batches:
            # recipe components set for this batch (identical across bottles)
            comps = [c for c,_ in self.bottles[bl[0]].recipe.components]
            seq = self.charging_sequence(comps)

            # batch-level precedence time: next component earliest start cannot be
            # earlier than previous component's end for THIS batch
            batch_prev_end = 0.0

            for comp in seq:
                # choose a mixer: earliest completion considering machine ready and batch precedence
                best = None
                # per-bottle filling time for this component:
                t_per_bottle = float(fill_t.get(comp, 0.02))
                # NOTE: we assume filling time not scaled by machine rate (you can change to: total_qty/rate)
                pt = t_per_bottle * len(bl)

                for m in range(self.D_mix):
                    stp = (0.0 if (last_comp[m] is None or last_comp[m] == comp)
                           else setup_diff)
                    if last_comp[m] is not None and last_comp[m] == comp:
                        stp = setup_same
                    st = max(ready[m] + stp, batch_prev_end)
                    en = st + pt
                    if (best is None) or (en < best[0]):
                        best = (en, m, st, pt, stp)
                en, m, st, pt_use, stp = best

                # record bar
                lbl = f"Batch({','.join('S'+str(i) for i in bl)})-{comp}"
                mix_tasks.append(MixTask(st, en, m, lbl))

                # update mixer state and batch precedence
                ready[m] = en
                last_comp[m] = comp
                batch_prev_end = en

            # after last component, all bottles in batch become mix_end = batch_prev_end
            for bid in bl:
                batch_last_end[bid] = batch_prev_end

        return batch_last_end, mix_tasks

    # ---------- QC helpers ----------
    def single_device_test_time(self, comp: str) -> float:
        return float(self.cfg["COMPONENT_TEST_TIME"].get(comp, 0.2))

    def choose_qc_mode(self, comps: List[str]) -> str:
        policy = self.cfg["QC_MODE_POLICY"]
        single = [self.single_device_test_time(c) for c in comps]
        mx, sm = max(single), sum(single)
        if policy == "always_1pass":
            return "1pass"
        if policy == "always_serial":
            return "serial"
        if policy == "by_ratio":
            return "1pass" if random.random() < self.cfg["ONEPASS_RANDOM_RATIO"] else "serial"
        return "1pass" if mx <= self.cfg["ONEPASS_RATIO"] * sm else "serial"

    # ---------- QC scheduling ----------
    def schedule_qc(self, mix_end: Dict[int,float]) -> Tuple[float, List[QCTask]]:
        D = self.D_qc
        dev_comp = self.dev_to_comp
        ready = [0.0]*D
        last_sig: List[Optional[str]] = [None]*D
        setup_t = self.cfg["QC_SETUP_TIME"]

        # bottles in order of ready time from MIX
        order_seq = sorted(self.bottles, key=lambda b: mix_end[b.id])
        tasks: List[QCTask] = []

        # map component -> list of candidate devices (single type per device)
        comp_to_devs: Dict[str, List[int]] = {}
        for d, c in enumerate(dev_comp):
            comp_to_devs.setdefault(c, []).append(d)

        for b in order_seq:
            bid = b.id
            comps = [c for c,_ in b.recipe.components]

            # pick a device for each required component (choose earliest ready+setup)
            picked_devs: List[int] = []
            for comp in comps:
                cands = comp_to_devs.get(comp, [])
                if not cands:
                    raise RuntimeError(f"No QC device can test component {comp}")
                best_d = None; best_ready = None
                for d in cands:
                    stp = 0.0 if (last_sig[d] is None or last_sig[d]==comp) else setup_t
                    rdy = ready[d] + stp
                    if (best_ready is None) or (rdy < best_ready):
                        best_ready = rdy; best_d = d
                picked_devs.append(best_d)

            mode = self.choose_qc_mode(comps)

            if mode == "1pass":
                # all devices must start together at the max of (device ready+setup) and mix_end
                earliest = mix_end[bid]
                per = []
                for d, comp in zip(picked_devs, comps):
                    stp = 0.0 if (last_sig[d] is None or last_sig[d]==comp) else setup_t
                    earliest = max(earliest, ready[d] + stp)
                    pt = self.single_device_test_time(comp)
                    per.append((d, comp, stp, pt))
                seg_len = max(pt for (_,_,_,pt) in per)
                st = earliest; en = st + seg_len
                group_key = (bid, round(st,6), round(en,6))
                used = "+".join(f"E{d+1}" for (d,_,_,_) in per)
                label = f"S{bid} 1-pass [{used}] max={seg_len:.2f}h"
                for d, comp, stp, pt in per:
                    tasks.append(QCTask(st, en, d, label, group_key))
                    ready[d] = en
                    last_sig[d] = comp
            else:
                # serial by device id
                cur = mix_end[bid]
                for d, comp in sorted(zip(picked_devs, comps), key=lambda x: x[0]):
                    stp = 0.0 if (last_sig[d] is None or last_sig[d]==comp) else setup_t
                    st = max(cur, ready[d] + stp)
                    pt = self.single_device_test_time(comp)
                    en = st + pt
                    tasks.append(QCTask(st, en, d, f"S{bid} serial E{d+1}({comp})", None))
                    ready[d] = en
                    last_sig[d] = comp
                    cur = en

        Cmax = max((t.end for t in tasks), default=0.0)
        return Cmax, tasks


# ============================ Plotting ============================ #

def wrap_text(s: str, width: int) -> str:
    # very simple wrap by characters; break at spaces if possible
    out=[]; line=""
    for ch in s:
        line += ch
        if len(line) >= width and ch == ' ':
            out.append(line.rstrip()); line=""
    if line: out.append(line)
    return "\n".join(out)

def plot_gantt(mix_tasks: List[MixTask], qc_tasks: List[QCTask], cfg,
               title="Mix & QC Gantt (component-wise MIX; 1-pass same on all devices)"):
    if not HAS_MPL:
        print("(matplotlib not installed; skip plotting)"); return

    # rows: MIX machines then QC devices
    rows=[]
    for m in range(cfg["NUM_MIX_MACHINES"]):
        rows.append(("MIX", m))
    for d in range(cfg["NUM_QC_DEVICES"]):
        rows.append(("QC", d))
    row_idx = {("MIX",m): m for m in range(cfg["NUM_MIX_MACHINES"])}
    off = cfg["NUM_MIX_MACHINES"]
    for d in range(cfg["NUM_QC_DEVICES"]):
        row_idx[("QC",d)] = off + d

    # colors for 1-pass groups
    from itertools import cycle
    palette = cycle(plt.rcParams['axes.prop_cycle'].by_key()['color'])
    grp_color: Dict[Tuple[int,float,float], str] = {}
    for t in qc_tasks:
        if t.group_key and t.group_key not in grp_color:
            grp_color[t.group_key] = next(palette)

    fig_h = max(4.0, 0.6*len(rows)+2)
    fig, ax = plt.subplots(figsize=(12, fig_h))

    # draw MIX
    for t in mix_tasks:
        y = row_idx[("MIX", t.machine)]
        ax.barh(y, t.end-t.start, left=t.start, height=0.8)
        ax.text((t.start+t.end)/2, y,
                wrap_text(t.label, cfg["LABEL_WRAP_CHARS"]),
                ha='center', va='center', color='white', fontsize=8)

    # draw QC
    for t in qc_tasks:
        y = row_idx[("QC", t.device)]
        color = grp_color.get(t.group_key, None)
        ax.barh(y, t.end-t.start, left=t.start, height=0.8, color=color)
        ax.text((t.start+t.end)/2, y,
                wrap_text(t.label, cfg["LABEL_WRAP_CHARS"]),
                ha='center', va='center',
                color=('white' if color else 'black'), fontsize=8)

    # y labels
    yticks=[]; ylabels=[]
    for m in range(cfg["NUM_MIX_MACHINES"]):
        yticks.append(row_idx[("MIX", m)])
        ylabels.append(f"MIX-M{m+1}")
    # show device type on QC rows
    types = CONFIG["COMPONENT_TYPES"]
    for d in range(cfg["NUM_QC_DEVICES"]):
        yticks.append(row_idx[("QC", d)])
        ylabels.append(f"QC-E{d+1} ({types[d % len(types)]})")

    ax.set_yticks(yticks); ax.set_yticklabels(ylabels)
    ax.set_xlabel("Time (h)"); ax.set_title(title)
    ax.grid(axis='x', alpha=0.3)
    plt.tight_layout()
    plt.show()


# ============================ Main ============================ #

def main():
    bottles, dev_to_comp = RandomData(CONFIG).gen()

    sch = Scheduler(CONFIG, bottles, dev_to_comp)
    mix_end, mix_tasks = sch.schedule_mixing()               # 充气（逐组分）
    Cmax, qc_tasks = sch.schedule_qc(mix_end)                # 检查（1-pass/serial）

    print(f"Bottles: {CONFIG['NUM_BOTTLES']}, Mixers: {CONFIG['NUM_MIX_MACHINES']}, QC devices: {CONFIG['NUM_QC_DEVICES']}")
    print(f"QC Cmax: {Cmax:.3f} h")
    if HAS_MPL:
        plot_gantt(mix_tasks, qc_tasks, CONFIG)
    else:
        print("(Install matplotlib to view Gantt.)")

if __name__ == "__main__":
    main()
